import time
import traceback
from typing import Callable, Awaitable
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import StreamingResponse
from app.utils.logger import logger
from app.config.settings import settings

class ResponseCaptureMiddleware(BaseHTTPMiddleware):
    """
    响应捕获中间件，用于记录API响应体
    """

    async def dispatch(
        self,
        request: Request,
        call_next: Callable[[Request], Awaitable[Response]]
    ) -> Response:
        # 记录请求开始时间
        start_time = time.time()

        # 优先从请求头获取，其次从state获取
        request_id = (request.headers.get('X-Request-Id') or
                      getattr(request.state, 'request_id', 'N/A'))
        
        # 获取请求信息
        request_method = request.method
        request_url = str(request.url)
        client_ip = request.client.host if request.client else "unknown"
        
        # 尝试读取请求体
        request_body = None
        if request_method in ["POST", "PUT", "PATCH"]:
            try:
                body_bytes = await request.body()
                if body_bytes:
                    request_body = body_bytes.decode("utf-8")
            except Exception as e:
                logger.warning(f"[Request-ID: {request_id}] Failed to read request body: {e}")
        
        # 记录请求信息（包含请求体）
        request_log = f"[Request-ID: {request_id}] Request started - Method: {request_method}, URL: {request_url}, Client IP: {client_ip}"
        if request_body and settings.LOG_REQUEST_BODY:
            request_log += f", Body: {request_body}"
        logger.info(request_log)
        
        try:
            # 调用下一个中间件或路由处理函数
            response = await call_next(request)

            # 计算处理时间
            process_time = time.time() - start_time
            
            # 记录响应信息
            response_log = (f"[Request-ID: {request_id}] Request completed - "
                           f"Method: {request_method}, "
                           f"URL: {request_url}, "
                           f"Status: {response.status_code}, "
                           f"Process Time: {process_time:.4f}s")
            
            if settings.LOG_RESPONSE_BODY:
                # 尝试读取响应体
                if isinstance(response, StreamingResponse):
                    response_log += f", Body: [Streaming Response]"
                elif hasattr(response, 'body') and response.body:
                    try:
                        response_body = response.body.decode('utf-8')
                        response_log += f", Body: {response_body}"
                    except Exception as e:
                        logger.warning(f"[Request-ID: {request_id}] Failed to log response body: {e}")
                else:
                    # 对于其他类型的响应，记录基本信息
                    content_type = response.headers.get('content-type', '')
                    if 'application/json' in content_type:
                        response_log += f", Body: [JSON Response - Content not captured in middleware]"
                    else:
                        response_log += f", Content-Type: {content_type}"
            
            logger.info(response_log)
            
            # 确保响应头中包含X-Request-Id
            response.headers["X-Request-Id"] = request_id
            
            return response
            
        except Exception as e:
            # 记录异常
            process_time = time.time() - start_time
            error_log = (f"[Request-ID: {request_id}] Request failed - "
                         f"Method: {request_method}, "
                         f"URL: {request_url}, "
                         f"Error: {str(e)}, "
                         f"Process Time: {process_time:.4f}s")
            logger.error(error_log)
            logger.debug(f"[Request-ID: {request_id}] Traceback: {traceback.format_exc()}")
            
            # 即使在异常情况下也要确保响应头中包含X-Request-Id
            # 注意：在这种情况下我们不能修改响应头，因为异常可能发生在响应创建之后
            raise e

# 为了向后兼容，保留原来的类名
LoggingMiddleware = ResponseCaptureMiddleware
